import torch
from torch import Tensor
from torch.autograd import Variable
from torch.nn import Module
from torch.utils.data import DataLoader

from .TorchDataset import MyDataset


def train(model: Module, x: Tensor, y: Tensor, save_path: str):
    dataset = MyDataset(x, y)
    data_loader = DataLoader(dataset=dataset, batch_size=32, shuffle=True)
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    for index in range(1000):
        for x, y in data_loader:
            inputs = Variable(x)
            target = Variable(y)
            predict = model(inputs)
            loss = criterion(predict, target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print(loss.item())

    torch.save(model, save_path)
